/*
 *  findBallState.cpp
 *  Ball and Plate
 *
 *  Created by Kenn Sebesta on 2007-11-23.
 *  Copyright 2007 __MyCompanyName__. All rights reserved.
 *
 */

#include "findBallState.h"
#include "maskGenerate.h"
#include "displayOpenCV.h"

CvMat *convMat; //This was declared global because I couldn't figure out any other way to kill the memory leak.
                //IT NEEDS TO BE FIXED.

extern double convThresh;
extern int viewSize;
extern float decimate; //Only useful for turning off decimation in case of problems with the function.

void ballFinder(/*Input variables*/	IplImage *maskOriginal, IplImage *camFrame, ImageSet images, CvMat **dft_mIm, CvRect cropRect, ImageAnalysisParams params,
        /*Output variables*/    int *missingCount, double *convMin, double *convMax, CvMat *y_measrd, CvMat **convMatOUT, IplImage **edgeImageOUT) {
    
//cvWaitKey(0);
    // Intialize image matrices
//    IplImage *cropImage, *grayImage, *bwImage;
    IplImage *edgeImage, *maskImage;
    
    
    //Manipulate images
//    cvSetImageROI( camFrame, cropRect );
    
#if defined DEBUG && (DEBUG > 3)
    printf("Error 10\n");
#endif
    
//    if(!decimate){
//        cropImage = cvCreateImage(cvGetSize(camFrame), camFrame->depth, camFrame->nChannels);
//        maskImage = cvCreateImage(cvGetSize(maskOriginal), maskOriginal->depth, maskOriginal->nChannels);
//    }
//    else{
//        cropImage = cvCreateImage(cvSize((int)(viewSize/decimate), (int)(viewSize/decimate)), camFrame->depth, camFrame->nChannels);
    maskImage = cvCreateImage(cvSize((int)(maskOriginal->width*viewSize/decimate/(double)cropRect.width), (int)(maskOriginal->height*viewSize/decimate/(double)cropRect.height)), maskOriginal->depth, maskOriginal->nChannels);
//    }
    
//    grayImage = cvCreateImage(cvGetSize(cropImage), cropImage->depth, 1);
//    bwImage = cvCreateImage(cvGetSize(cropImage), cropImage->depth, 1);
    edgeImage = cvCreateImage(cvSize((int)(viewSize/decimate), (int)(viewSize/decimate)), images.edgeImage->depth, 1);
    
#if defined DEBUG && (DEBUG > 3)
    printf("Error 11\n");
#endif
    if((decimate<=1) && *missingCount<=0){
//        cvCopy( camFrame, cropImage );
        cvCopy( images.edgeImage, edgeImage );
        cvCopy( maskOriginal, maskImage );
    }
    else{
//        cvResize( camFrame, cropImage, CV_INTER_AREA );
        cvResize(images.edgeImage, edgeImage, CV_INTER_AREA);
        cvResize( maskOriginal, maskImage, CV_INTER_AREA );
    }
#if defined DEBUG && (DEBUG > 3)
    printf("Error 12\n");
#endif
    
//    if (cropImage->nChannels!=1){
//        cvSplit( cropImage, grayImage, NULL, NULL, NULL );
//        //	cvCvtColor(cropImage, grayImage, CV_BGR2GRAY);	//Can't figure out how to tell if it's a RGB or
//        //grayscale because my camera returns three
//        //channels, even though it's grayscale.
//    }
//    else{
//        cvCopy(cropImage, grayImage);
//    }
#if defined DEBUG && (DEBUG > 3)
    printf("Error 13\n");
#endif
    
#if defined DEBUG && (DEBUG > 3)
    printf("Error 20\n");
#endif
    
    //Image analysis
    /*	if(0)
     * cvThreshold(grayImage, bwImage, (double)binaryThresh, 255, CV_THRESH_BINARY); // Binarises the image. IS DEACTIVATED UNTIL I SET UP THE CAMERA OVER THE TABLE.
     * else
     * cvCopy( grayImage, bwImage );*/
//    cvCopy( grayImage, bwImage );
//    cvCanny( bwImage, edgeImage, params.cannyThreshLow, params.cannyThreshHigh, 3 ); // Performs the edge extraction on the image
//cvShowImage("Parameters", edgeImage);
//cvWaitKey(0);
//cvShowImage("Parameters", edgeImage);
//cvWaitKey(0);
    
    int dft_rows = cvGetOptimalDFTSize( edgeImage->height + maskImage->height - 1 );	//This code could be done offline, but there's no reason to because it's so fast.
    int dft_cols = cvGetOptimalDFTSize( edgeImage->width + maskImage->width - 1 );		//This code could be done offline, but there's no reason to because it's so fast.
    
    CvMat *tmp; //THIS IS POOR IMPLEMENTATION. convMAT SHOULD not BE GLOBAL.
    
#if defined DEBUG && (DEBUG > 3)
    printf("Error 30\n");
#endif
    
    maskDFT( dft_rows, dft_cols, *missingCount, maskImage, dft_mIm );
    
#if defined DEBUG && (DEBUG > 3)
    printf("Error 40\n");
#endif
    
    conv2D( /*Input variables*/     edgeImage, maskImage, *dft_mIm, dft_rows, dft_cols,
            /*Output variables*/    &tmp);
    cvReleaseMat(&tmp);
    //================================================
    // Returns the maximum convolution point, along with the max/min convolution values
    // in order to properly scale the output image.
    //================================================
    
#if defined DEBUG && (DEBUG > 3)
    printf("Error 50\n");
#endif
    
    CvPoint ballPoint;
    cvMinMaxLoc( convMat, convMin, convMax, NULL, &ballPoint);
    if ( *convMax > convThresh*pow(maskImage->width/(double)maskOriginal->width, 1.4) ){
        // The IF clause is complex because the convolution doesn't change linearly with mask size.
        // I'm open to all ideas for improving the IF clause.
        
        // Reset the missing count
        if ( *missingCount>0 )
            *missingCount=-1;
        else
            *missingCount=0;
        
        // Find the center of the ball to sub-pixel resolution
        centerOfMoments(edgeImage, convMat, ballPoint, cropRect, &y_measrd);
    }
    else{
        printf("convMax = %4.3e, convThresh = %4.3e\n",  *convMax, convThresh*pow(maskImage->width/(double)maskOriginal->width, 1.4));
        *missingCount=*missingCount+1;
    }
    //Outputs
//    images.convMat=cvCloneMat(convMat);
    *convMatOUT=cvCloneMat(convMat);
    *edgeImageOUT= cvCloneImage(edgeImage);
    
    // Garbage Collection
    cvResetImageROI( camFrame );
//	free( ballPoint );
//    cvReleaseImage( &cropImage );
//    cvReleaseImage( &grayImage );
//    cvReleaseImage( &bwImage );
    cvReleaseImage( &maskImage );
    cvReleaseImage( &edgeImage );
    cvReleaseMat( &convMat );
    
//    delete &ballPoint;
    return;
}


//============================================================================================

void conv2D(/*Input variables*/     IplImage *edgeImage, IplImage *maskImage, CvMat *dft_mIm, int dft_rows, int dft_cols,
        /*Output variables*/    CvMat **convMatOUT){
    
// This function takes two images and computes the convolution. It was
    // copied directly off of the OpenCV wiki. I made only minor modifications
    // to adjust for variable types.
    
//The convolution matrix is convMat, which is global. THIS IS POOR PROGRAMMING. FIX IT!
    
    // Declare structures
    IplImage *edgeImageF;
    CvMat *conv, *dft_eIm;
    CvMat tmp; //Temporary matrix which is NECESSARY. Won't work without it.
    conv = cvCreateMat( edgeImage->height + maskImage->height - 1, edgeImage->width + maskImage->width - 1, CV_32F );
    dft_eIm = cvCreateMat( dft_rows, dft_cols, conv->type );
    edgeImageF= cvCreateImage(cvGetSize(edgeImage),  IPL_DEPTH_32F, 1); //Edge matrix in F32
    
    // Create image structures
    
#if defined DEBUG && (DEBUG > 3)
    printf("Error 41\n");
    printf("dft_eIm = %d, %d\n", dft_eIm->rows, dft_eIm->rows );
    printf("edgeImage = %d, %d\n", edgeImage->height, edgeImage->width);
    printf("edgeImageF = %d, %d\n", edgeImageF->height, edgeImageF->width);
#endif
    
    cvConvert( edgeImage, edgeImageF ); //Change from integer to float
    
#if defined DEBUG && (DEBUG > 3)
    printf("Error 42\n");
#endif    
    // Copy edgeImageF to dft_eIm and pad dft_eIm with zeros
    cvGetSubRect( dft_eIm, &tmp, cvRect( 0, 0, edgeImage->width, edgeImage->height ));
    cvCopy( edgeImageF, &tmp );
    cvGetSubRect( dft_eIm, &tmp, cvRect( edgeImage->width, 0, dft_cols - edgeImage->width, edgeImage->height ));
    cvZero( &tmp );
    
#if defined DEBUG && (DEBUG > 3)
    printf("Error 43\n");
#endif
    
    // No need to pad bottom part of dft_eIm with zeros because of use of non-zero rows parameter in cvDFT() calls below
    cvDFT( dft_eIm, dft_eIm, CV_DXT_FORWARD, edgeImage->height );
    
#if defined DEBUG && (DEBUG > 3)
    printf("Error 46\n");
//    printf("dft_eIm %d, %d\n", dft_eIm->rows, dft_eIm->cols);
    printf("dft_mIm %d, %d\n", dft_mIm->rows, dft_mIm->cols);
#endif
    
    
    //Perform the actual convolution
    cvMulSpectrums( dft_eIm, dft_mIm, dft_eIm, 0 );
    
#if defined DEBUG && (DEBUG > 3)
    printf("Error 47\n");
#endif
    
    //Inverses the DFT into a regular image
    cvDFT( dft_eIm, dft_eIm, CV_DXT_INV_SCALE, conv->rows ); // calculate only the top part
    
    //Get only the region that has data. This should return a matrix of the same size and
    //position as the input image.
    cvGetSubRect( dft_eIm, &tmp, cvRect( maskImage->width/2, maskImage->height/2, edgeImage->width, edgeImage->height ));//Might need fixing due to the x,y coordinates of the CvRect.
    
    //NECESSARY matrix clone. Won't work without it.
    convMat=cvCloneMat(&tmp);
    *convMatOUT=cvCloneMat(&tmp);
    
#if defined DEBUG && (DEBUG > 3)
    printf("Error 49\n");
#endif
    
    //Garbage Collection
    cvReleaseImage( &edgeImageF );
    cvReleaseMat( &dft_eIm );
    cvReleaseMat( &conv );
    
    return;
    
}

void centerOfMoments( IplImage *edgeImage, CvMat *convMat, CvPoint ballPoint, CvRect cropRect, CvMat **y_measrd)
// Finds the center of the convolution through the method of moments.
{
    int pixelNeighborhood=5, x0, x1, y0, y1;
    double weight;
    CvMat *xLeverArm, *yLeverArm, *xSumTmp, *ySumTmp;
    CvMat baryMat;
    
    // Calculation of a small window around ballPoint, with a verification that we don't leave the convolution image.
    x0=std::max( ballPoint.x - pixelNeighborhood, 0 );
    x1=std::min( ballPoint.x + pixelNeighborhood, edgeImage->width-1);
    y0=std::max( ballPoint.y - pixelNeighborhood, 0 );
    y1=std::min( ballPoint.y + pixelNeighborhood, edgeImage->height-1);
    xLeverArm=cvCreateMat( x1-x0+1, 1, CV_32F );
    yLeverArm=cvCreateMat( 1, y1-y0+1, CV_32F );
    xSumTmp=cvCreateMat( y1-y0+1, 1, CV_32F );
    ySumTmp=cvCreateMat( 1, x1-x0+1, CV_32F );
    for ( int i=x0; i<=x1 ; i++ )
        cvmSet( xLeverArm, i-x0, 0, (double)i );
    for ( int i=y0; i<=y1 ; i++ )
        cvmSet( yLeverArm, 0, i-y0, (double)i );
    
    cvGetSubRect( convMat, &baryMat, cvRect( x0, y0, x1-x0+1, y1-y0+1 ) ); //crops image
    
    // Normalization...
    weight=cvSum( &baryMat ).val[0];
    
    // ... and barycenter:
    //	if( weight>1e02 && *convMax > 1.5e06 ){ /*Weight only equals 0 when we can't find the ball*/
    cvMatMul( &baryMat, xLeverArm, xSumTmp );
    cvMatMul( yLeverArm, &baryMat, ySumTmp );
    
    // Set coordinates relative to absolute origin.
    if (cropRect.width!=0){
        cvmSet( *y_measrd, 0, 0, cvSum( xSumTmp ).val[0]/weight*cropRect.width/viewSize*decimate+cropRect.x);
        cvmSet( *y_measrd, 1, 0, cvSum( ySumTmp ).val[0]/weight*cropRect.width/viewSize*decimate+cropRect.y);
    }
    else{
        cvmSet( *y_measrd, 0, 0, cvSum( xSumTmp ).val[0]/weight+cropRect.x);
        cvmSet( *y_measrd, 1, 0, cvSum( ySumTmp ).val[0]/weight+cropRect.y);
    }

    //		printf("x=%f, y=%f absolute\n", cvmGet( y_measrd, 0, 0 ), cvmGet( y_measrd, 1, 0 ));
    
    //Local garbage collection
    cvReleaseMat( &xSumTmp );
    cvReleaseMat( &ySumTmp );
    cvReleaseMat( &xLeverArm );
    cvReleaseMat( &yLeverArm );
    
    return;
}
